import os
import re
import shutil
import time
from datetime import datetime

import pandas as pd


def clean_pti_ssdd(input_folder, output_folder):
    """
    处理Excel文件：
    1. SSDD工作表按原逻辑处理
    2. OPO工作表（不区分大小写）直接提取原始内容，按最新日期保存并替换旧文件
    3. 增加文件占用容错处理
    """
    source_dir = os.path.join(input_folder, "PTI")
    target_dir = os.path.join(output_folder, "PTI")
    archive_dir = os.path.join(input_folder, "Archive", "PTI")
    opo_target_dir = os.path.join(output_folder, "OPO", "PTI")
    # 检查目录有效性
    if not os.path.exists(source_dir):
        print(f"错误: 源目录 '{source_dir}' 不存在")
        return
    os.makedirs(target_dir, exist_ok=True)
    os.makedirs(archive_dir, exist_ok=True)
    os.makedirs(opo_target_dir, exist_ok=True)
    print(f"SSDD处理后保存至: {target_dir}")
    print(f"源文件归档至: {archive_dir}")
    print(f"OPO提取后保存至: {opo_target_dir}")

    # 收集Excel文件
    excel_exts = ('.xlsx', '.xls', '.xlsm')
    excel_files = [f for f in os.listdir(source_dir) if f.lower().endswith(excel_exts)]
    if not excel_files:
        print(f"源目录 '{source_dir}' 中无Excel文件")
        return

    # 存储OPO文件信息（路径+日期+数据）
    opo_file_info = []

    for file in excel_files:
        src_path = os.path.join(source_dir, file)
        max_retries = 3
        retry_count = 0
        while retry_count < max_retries:
            try:
                # 读取Excel文件
                excel_file = pd.ExcelFile(src_path)
                sheet_names = excel_file.sheet_names

                # 处理SSDD工作表（保留原逻辑）
                if 'SSDD' in sheet_names:
                    df = excel_file.parse('SSDD', header=None)
                    if len(df) < 3:
                        print(f"警告: {file} 的SSDD表行数不足，跳过处理")
                    else:
                        df = df.drop([0, 1]).reset_index(drop=True)
                        df.columns = df.iloc[0]
                        df = df.drop(0).reset_index(drop=True)

                        if len(df.columns) <= 19:
                            print(f"警告: {file} 的SSDD表无T列及以后数据，跳过处理")
                        else:
                            t_before = df.columns[:19]
                            t_after = df.columns[19:]
                            pivot_cols = [col for col in t_after if '~' in str(col) and '/' in str(col)]
                            keep_cols = [col for col in t_after if col not in pivot_cols]
                            id_cols = list(t_before) + keep_cols

                            if pivot_cols:
                                melted_df = pd.melt(df, id_vars=id_cols, value_vars=pivot_cols,
                                                    var_name='Original_Column', value_name='Value')
                                melted_df['Value'] = pd.to_numeric(melted_df['Value'], errors='coerce').fillna(0)
                                output_df = melted_df
                            else:
                                output_df = df

                            base_name = os.path.splitext(file)[0]
                            target_path = os.path.join(target_dir, f"{base_name}_revised.xlsx")
                            output_df.to_excel(target_path, index=False, engine='xlsxwriter')
                            print(f"SSDD处理完成: {target_path}")

                # 处理OPO工作表（不区分大小写，支持OPO、opo等形式）
                # 查找所有大小写形式匹配的工作表
                opo_sheets = [sheet for sheet in sheet_names if sheet.lower() == 'opo']
                if opo_sheets:
                    # 取第一个匹配的工作表（通常只有一个）
                    opo_sheet_name = opo_sheets[0]
                    # 直接读取原始内容（保留所有行和列，不修改表头）
                    df_opo = excel_file.parse(opo_sheet_name)

                    # 解析文件名中的日期（匹配YYYYMMDD格式）
                    date_match = re.search(r'(\d{8})', file)
                    file_date = datetime.min  # 默认最小日期
                    if date_match:
                        try:
                            file_date = datetime.strptime(date_match.group(1), '%Y%m%d')
                        except:
                            print(f"警告: {file} 文件名日期格式不正确，按默认日期处理")

                    opo_file_info.append({
                        'filename': file,
                        'date': file_date,
                        'data': df_opo
                    })
                    print(f"已识别并提取 {file} 中的 {opo_sheet_name} 工作表")

                # 显式关闭ExcelFile以释放资源
                excel_file.close()

                # 归档源文件（增加重试逻辑）
                archive_path = os.path.join(archive_dir, file)
                shutil.move(src_path, archive_path)
                print(f"源文件已归档: {archive_path}")
                break  # 处理成功，跳出重试循环

            except PermissionError as e:
                if "being used by another process" in str(e):
                    retry_count += 1
                    print(f"文件被占用，第 {retry_count} 次重试...")
                    time.sleep(1)  # 等待1秒后重试
                else:
                    print(f"处理 {file} 出错: {str(e)}")
                    break
            except Exception as e:
                print(f"处理 {file} 出错: {str(e)}")
                break

    # 保存最新的OPO文件（替换旧文件）
    if opo_file_info:
        # 按日期降序排序，取最新的一个
        opo_file_info.sort(key=lambda x: x['date'], reverse=True)
        latest_opo = opo_file_info[0]

        # 生成目标文件名（保持原命名规则）
        base_name = os.path.splitext(latest_opo['filename'])[0]
        opo_target_path = os.path.join(opo_target_dir, f"{base_name}_opo.xlsx")

        # 保存原始OPO内容（不修改索引）
        latest_opo['data'].to_excel(opo_target_path, index=False, engine='openpyxl')
        print(f"最新OPO文件已保存（替换旧文件）: {opo_target_path}")
    else:
        print("无OPO工作表需处理")


def clean_pega_ssdd(input_folder, output_folder):
    print("Cleaning PEGA SSDD files...")
    target_dir = os.path.join(output_folder, "Pega")
    archive_dir = os.path.join(input_folder, "Archive", "Pega")
    source_dir = os.path.join(input_folder, "Pega")
    os.makedirs(target_dir, exist_ok=True)  # 确保目标目录存在
    os.makedirs(archive_dir, exist_ok=True)  # 确保归档目录存在

    # 收集需处理的文件（避免边遍历边移动导致的列表变化）
    process_files = [
        f for f in os.listdir(source_dir)
        if f.endswith(('.xlsx', '.xls')) and "SSDD" in f
    ]

    for filename in process_files:
        src_path = os.path.join(source_dir, filename)
        print(f"开始处理: {filename}")
        try:
            df = pd.read_excel(src_path)

            ### 步骤4：删除“Group ID”列之前的所有列 ###
            if "Group ID" in df.columns:
                group_id_idx = df.columns.get_loc("Group ID")
                df = df.iloc[:, group_id_idx:]
            else:
                print(f"警告: {filename} 无'Group ID'列")
                # 无Group ID列时跳过后续依赖Group ID的步骤
                continue

            ### 步骤5：删除列名伪重复和内容重复的列 ###
            # 处理列名后缀（如 IPN.1 → IPN）
            base_cols = [re.sub(r'\.\d+$', '', str(col)) for col in df.columns]
            # 保留基础列名首次出现的列
            unique_cols = []
            seen_base = set()
            for col, base in zip(df.columns, base_cols):
                if base not in seen_base:
                    seen_base.add(base)
                    unique_cols.append(col)
            df = df[unique_cols]
            # 删除内容完全重复的列
            final_cols = []
            for col in df.columns:
                is_dup = any(df[col].equals(df[kept]) for kept in final_cols)
                if not is_dup:
                    final_cols.append(col)
            df = df[final_cols]

            ### 步骤6：仅保留含V/D的未命名列，其余删除 + 同一Group ID内有V则全部替换为V ###
            unnamed_cols = [col for col in df.columns if str(col).startswith("Unnamed:")]
            if unnamed_cols:
                # 筛选含V/D的未命名列（忽略空值）
                vd_unnamed = [
                    col for col in unnamed_cols
                    if df[col].astype(str).str.contains(r'[VD]', na=False).any()
                ]
                # 删除其他未命名列
                to_drop = [col for col in unnamed_cols if col not in vd_unnamed]
                df = df.drop(columns=to_drop, errors='ignore')

                # 命名首个含V/D的未命名列并处理值
                if vd_unnamed:
                    vmi_col = "VMI included in Group "
                    df = df.rename(columns={vd_unnamed[0]: vmi_col})

                    # 清洗列值：仅保留V/D，其余设为空白；处理特殊空值
                    df[vmi_col] = df[vmi_col].astype(str).str.strip().str.upper()
                    # 非V/D设为空白，特殊值（NONE、NaN等）也设为空白
                    df[vmi_col] = df[vmi_col].where(df[vmi_col].isin(['V', 'D']), '').replace(['NONE', 'NaN', 'nan'],
                                                                                              '')

                    # 自定义函数：确保返回同长度Series（解决拼接报错）
                    def replace_with_v(series):
                        if (series == "V").any():
                            return pd.Series(["V"] * len(series), index=series.index)
                        return series

                    # 按Group ID分组替换
                    df[vmi_col] = df.groupby("Group ID")[vmi_col].apply(replace_with_v).reset_index(level=0, drop=True)

            ### 步骤7：按Group ID填充IPN空值 ###
            if "IPN" in df.columns:
                df["IPN"] = df.groupby("Group ID")["IPN"].transform(
                    lambda x: x.fillna(x.dropna().iloc[0] if not x.dropna().empty else None)
                )
            else:
                print(f"警告: {filename} 无'IPN'列")

            ### 步骤8：日期列逆透视 ###
            date_pattern = re.compile(r'(\d{4}[-/]\d{1,2}[-/]\d{1,2})|(\d{1,2}[-/]\d{1,2}[-/]\d{4})')
            date_cols = [col for col in df.columns if date_pattern.search(str(col))]
            id_vars = [col for col in df.columns if col not in date_cols]
            if date_cols:
                df = pd.melt(df, id_vars=id_vars, value_vars=date_cols,
                             var_name="Date", value_name="Value")

            ### 步骤9：Value列转数值并填充0 ###
            if "Value" in df.columns:
                df["Value"] = pd.to_numeric(df["Value"], errors="coerce").fillna(0)

            ### 保存处理后文件 ###
            name, ext = os.path.splitext(filename)
            target_path = os.path.join(target_dir, f"{name}_Revised{ext}")
            df.to_excel(target_path, index=False, engine='xlsxwriter')  # 指定引擎避免格式问题
            print(f"处理后保存至: {target_path}")

            ### 归档源文件 ###
            archive_path = os.path.join(archive_dir, filename)
            shutil.move(src_path, archive_path)
            print(f"源文件归档至: {archive_path}")

        except Exception as e:
            print(f"处理 {filename} 出错: {str(e)}")
            # 出错时仍尝试归档源文件（避免残留）
            try:
                archive_path = os.path.join(archive_dir, filename)
                shutil.move(src_path, archive_path)
                print(f"出错文件已归档: {archive_path}")
            except:
                print(f"归档 {filename} 失败")



if __name__ == "__main__":
    input_folder = "test"
    output_folder = "output"
    # clean_pti_ssdd(input_folder, output_folder)
    # clean_pega_ssdd(input_folder, output_folder)